import time
import xml.etree.ElementTree as ET
import os
from os import getcwd
from universal.root_path import scanAllImages
try:
    from imgaug.augmenters.geometric import Rot90
    import imgaug as ia
    from imgaug import augmenters as iaa
    ia.seed(1)
except Exception as e:
    print(e)

import numpy as np
from PIL import Image

import shutil

# 数据增强脚本


# isOK=False

def read_xml_annotation(xml_path):
    in_file = open(xml_path, encoding='utf-8')
    tree = ET.parse(in_file)
    root = tree.getroot()
    bndboxlist = []
    classes = []
    for object in root.findall('object'):
        bndbox = object.find('bndbox')
        cls = object.find('name').text
        classes.append(cls)
        xmin = int(bndbox.find('xmin').text)
        xmax = int(bndbox.find('xmax').text)
        ymin = int(bndbox.find('ymin').text)
        ymax = int(bndbox.find('ymax').text)

        bndboxlist.append([xmin, ymin, xmax, ymax])

    return classes, bndboxlist


def change_xml_annotation(root, image_id, new_target):
    new_xmin = new_target[0]
    new_ymin = new_target[1]
    new_xmax = new_target[2]
    new_ymax = new_target[3]

    in_file = open(os.path.join(root, str(image_id) + '.xml'))
    tree = ET.parse(in_file)
    xmlroot = tree.getroot()
    object = xmlroot.find('object')
    bndbox = object.find('bndbox')
    xmin = bndbox.find('xmin')
    xmin.text = str(new_xmin)
    ymin = bndbox.find('ymin')
    ymin.text = str(new_ymin)
    xmax = bndbox.find('xmax')
    xmax.text = str(new_xmax)
    ymax = bndbox.find('ymax')
    ymax.text = str(new_ymax)
    tree.write(os.path.join(root, str(image_id) + "_aug" + '.xml'))


def change_xml_list_annotation(root, image_id, new_target, saveroot, id):

    in_file = open(os.path.join(root, str(image_id) + '.xml'), encoding='utf-8')
    tree = ET.parse(in_file)
    xmlroot = tree.getroot()
    index = 0

    for object in xmlroot.findall('object'):
        cls = object.find('name')
        bndbox = object.find('bndbox')

        new_xmin = new_target[index][0]
        new_ymin = new_target[index][1]
        new_xmax = new_target[index][2]
        new_ymax = new_target[index][3]

        xmin = bndbox.find('xmin')
        xmin.text = str(new_xmin)
        ymin = bndbox.find('ymin')
        ymin.text = str(new_ymin)
        xmax = bndbox.find('xmax')
        xmax.text = str(new_xmax)
        ymax = bndbox.find('ymax')
        ymax.text = str(new_ymax)

        index = index + 1

    tree.write(os.path.join(saveroot, str(image_id) + "_aug_" + str(id) + '.xml'))
    print(str(image_id[:-4]) + "_aug_" + str(id) + '.jpg')


def mkdir(path):
    path = path.strip()
    path = path.rstrip("\\")

    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        os.makedirs(path)
        print(path + ' 创建成功')
        return True
    # else:
    #     print(path + ' 目录已存在')
    #     return False


def StartIMG(cfg):
# if __name__ == "__main__":
    boxes_img_aug_list = []
    new_bndbox = []
    new_bndbox_list = []

    # AUGLOOP = 2  # 每张影像增强的数量
    # IMG_DIR = r"C:\Users\SV00299057\Desktop\99999"
    # XML_DIR = r"C:\Users\SV00299057\Desktop\99999"
    # AUG_IMG_DIR = r"C:\Users\SV00299057\Desktop\99999"  # 存储增强后的影像文件夹路径
    # AUG_XML_DIR = r"C:\Users\SV00299057\Desktop\99999"  # 存储增强后的XML文件夹路径

    AUGLOOP = int(cfg['EnhanceTimes'])  # 每张影像增强的数量
    IMG_DIR = cfg['InitImgPath']  #原路径
    XML_DIR = IMG_DIR

    AUG_IMG_DIR = cfg['SaveImgPath']  # 储存路径
    AUG_XML_DIR = AUG_IMG_DIR

    # try:
    #     shutil.rmtree(AUG_XML_DIR)
    # except FileNotFoundError as e:
    #     a = 1
    mkdir(AUG_XML_DIR)

    # try:
    #     shutil.rmtree(AUG_IMG_DIR)
    # except FileNotFoundError as e:
    #     a = 1
    mkdir(AUG_IMG_DIR)
    kwargs = {

    }

    # 影像增强

    seq = iaa.Sequential([
        iaa.Fliplr(float(cfg['Fliplr'])),
        iaa.Multiply(eval(cfg['Multiply'])), #或者tuple([cfg['Multiply']]) # change brightness, doesn't affect BBs
        iaa.GaussianBlur(sigma=eval(cfg['GaussianBlur'])),
        iaa.Affine(
            translate_px={"x": int(cfg['translate_X']), "y": int(cfg['translate_Y'])},
            scale=eval(cfg['scale']),  # 图像缩放为80%到95%之间
            rotate=eval(cfg['rotate'])  # 旋转±30度之间
        )
    ])

    # seq = iaa.Sequential([
    #     iaa.Fliplr(0.5),
    #     iaa.Multiply((0.9, 1.1)),  # change brightness, doesn't affect BBs
    #     iaa.GaussianBlur(sigma=(0, 3.0)),
    #     iaa.Affine(
    #         translate_px={"x": 15, "y": 15},
    #         scale=(0.9, 1.1),  # 图像缩放为80%到95%之间
    #         rotate=(-10, 10)  # 旋转±30度之间
    #     )
    # ])

    count = 0
    # for root, sub_folders, files in os.walk(XML_DIR):
    imglist = scanAllImages(XML_DIR)
    for img_path in imglist:
        try:
            if cfg['main_queue'].get(block=False).get('stop', None):
                cfg['main_queue'].put({'stop': True})
                break
        except:
            pass

        xml = os.path.splitext(img_path)[0]+'.xml'
        name = os.path.basename(xml)
        XML_DIR_ = os.path.dirname(xml)
        if os.path.exists(xml):
            classes, bndbox = read_xml_annotation(xml)

        # for i, name in enumerate(files):
        #     if not name.endswith('xml'):
        #         continue
        #     classes, bndbox = read_xml_annotation(XML_DIR, name)

            # 只保留翻折与开叉数据，剔除normal等其他类别
            # print(classes.__len__())
            # print(classes[0])
            # if('battery' not in classes and 'fork' not in classes):
            # if ('battery' not in classes):
            #     # if(classes.__len__()==1 and classes[0]=='normal'):
            #     continue
            print('>>> ', img_path)

            for epoch in range(AUGLOOP):
                seq_det = seq.to_deterministic()

                # 读取图片
                img = Image.open(img_path)
                img = np.array(img)

                index = []
                # bndbox 坐标增强
                for i in range(len(bndbox)):
                    bbs = ia.BoundingBoxesOnImage([
                        ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]),
                    ], shape=img.shape)

                    bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]

                    boxes_img_aug_list.append(bbs_aug)

                    new_bndbox_list.append([int(bbs_aug.bounding_boxes[0].x1),
                                            int(bbs_aug.bounding_boxes[0].y1),
                                            int(bbs_aug.bounding_boxes[0].x2),
                                            int(bbs_aug.bounding_boxes[0].y2)])
                    bb = bbs.bounding_boxes[0]
                    # 判断是否处于图形区域内
                    if bb.is_fully_within_image(img.shape):
                        id = 1
                    elif bb.is_partly_within_image(img.shape):
                        id = 1
                    else:
                        id = 0

                    n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
                    n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
                    n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
                    n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
                    # 限制增强后区域位置（x1,y1;x2,y2）
                    if (n_x1 < 20 or n_x2 > (img.shape[1] - 20)) or abs(n_x1 - n_x2) < 15:
                        id = 0
                    if (n_y1 < 20 or n_y2 > (img.shape[0] - 20)) or abs(n_y1 - n_y2) < 15:
                        id = 0
                    if n_x1 >= n_x2 or n_y1 >= n_y2:
                        id = 0

                    index.append(id)

                if len(new_bndbox_list) == 0:
                    continue

                # 存储变化后的XML
                change_xml_list_annotation(XML_DIR_, name[:-4], new_bndbox_list, XML_DIR_, epoch)
                new_bndbox_list = []

                image_aug = seq_det.augment_images([img])[0]
                path = os.path.join(XML_DIR_, str(name[:-4]) + "_aug_" + str(epoch) + '.jpg')
                Image.fromarray(image_aug).save(path)

    cfg['main_queue'].put({'stop': True})


